#ifndef LONGKEY_PCAGENT_COM_PROXY_H_
#define LONGKEY_PCAGENT_COM_PROXY_H_

#include "com_proxy_clsid.h"
#include "longkey_pcagent_ps_resource.h"

// All[*] coclasses in pcagent_idl.idl do the following:
// * Derive from StdMarshalInfo.
// * Construct StdMarshalInfo(is_machine).
// * Add a COM_INTERFACE_ENTRY(IStdMarshalInfo)
//
// [*] The proxy classes LongkeyComProxy[XXX]Class, being proxies, do not follow
// the steps above. In addition, the CurrentStateClass has a custom marshaler
// and does not follow the steps above.
//
const IID kIIDsToRegister[] = 
{
	__uuidof(ILKPCAgent),
};

struct ComProxyMode {
	static bool is_machine() { return false; }

	static const GUID& class_id()
	{
		return is_machine() ? __uuidof(LongkeyComProxyMachineClass) :
			__uuidof(LongkeyComProxyUserClass);
	}

	static const GUID ps_clsid()
	{
		if (is_machine())
		{
			GUID proxy_clsid = PROXY_CLSID_IS_MACHINE;
			return proxy_clsid;
		} else {
			GUID proxy_clsid = PROXY_CLSID_IS_USER;
			return proxy_clsid;
		}
	}

	static const TCHAR* const hk_root()
	{
		return is_machine() ? _T("HKLM") : _T("HKCU");
	}
};

#pragma warning(push)

// C4640: construction of local static object is not thread-safe
#pragma warning(disable : 4640)

// C4505: unreferenced IUnknown local functions have been removed
#pragma warning(disable : 4505)

class ATL_NO_VTABLE ComProxy : public CComObjectRootEx<CComMultiThreadModel>,
	public CComCoClass<ComProxy>, public IUnknown
{
public:
	ComProxy() { CORE_LOG(L2, (_T("[ComProxy::ComProxy]"))); }

	DECLARE_GET_CONTROLLING_UNKNOWN()
	DECLARE_REGISTRY_RESOURCEID_EX(IDR_COM_PROXY_RGS);

#pragma warning(push)
	// Construction of local static object is not thread-safe
#pragma warning(disable:4640)
	BEGIN_REGISTRY_MAP()
		REGMAP_ENTRY(_T("HKROOT"), ComProxyMode::hk_root())
		REGMAP_ENTRY(_T("CLSID"),  ComProxyMode::class_id())
	END_REGISTRY_MAP()
#pragma warning(pop)

	BEGIN_COM_MAP(ComProxy)
		COM_INTERFACE_ENTRY(IUnknown)
		COM_INTERFACE_ENTRY_FUNC(__uuidof(IClientSecurity), 0, QueryInternal)
		COM_INTERFACE_ENTRY_FUNC(__uuidof(IMultiQI), 0, QueryInternal)
		COM_INTERFACE_ENTRY_AGGREGATE_BLIND(proxy_manager_.p)
	END_COM_MAP()

	static HRESULT WINAPI QueryInternal(void* ptr, REFIID iid,
	void** retval, DWORD_PTR) {
		ASSERT1(ptr);
		ASSERT1(retval);
		CORE_LOG(L2, (_T("[ComProxy::QueryInternal][%s]"), GuidToString(iid)));

		ComProxy* this_ptr = reinterpret_cast<ComProxy*>(ptr);
		return this_ptr->proxy_internal_unknown_->QueryInternalInterface(iid,
			retval);
	}

	HRESULT FinalConstruct() {
		CORE_LOG(L2, (_T("[ComProxy::FinalConstruct]")));

		HRESULT hr = RegisterProxyStubs();
		if (FAILED(hr)) {
			CORE_LOG(LE, (_T("[RegisterProxyStubs failed][0x%x]"), hr));
			// If explicit registration failed, the registry-based proxy lookup
			// mechanism may still work. Fall through.
		}

		hr = ::CoGetStdMarshalEx(GetControllingUnknown(), SMEXF_HANDLER, &proxy_manager_);
		if (FAILED(hr)) {
			CORE_LOG(LE, (_T("[::CoGetStdMarshalEx failed][0x%x]"), hr));
			return hr;
		}

		return proxy_manager_.QueryInterface(&proxy_internal_unknown_);
	}

	void FinalRelease()
	{
		CORE_LOG(L2, (_T("[ComProxy::FinalRelease]")));
	}

	static HRESULT RegisterProxyStubs()
	{
		static LLock lock;
		static bool is_registered = false;

		__mutexScope(lock);

		if (is_registered) 
			return S_OK;

		CORE_LOG(L2, (_T("[ComProxy::RegisterProxyStubs][Registering][%d]"),
			ComProxyMode::is_machine()));

		const GUID ps_clsid = ComProxyMode::ps_clsid();
		for (size_t i = 0; i < arraysize(kIIDsToRegister); ++i)
		{
			HRESULT hr = ::CoRegisterPSClsid(kIIDsToRegister[i], ps_clsid);
			if (FAILED(hr))
			{
				CORE_LOG(LE, (_T("[::CoRegisterPSClsid failed][%s][%s][0x%x]"),
					GuidToString(kIIDsToRegister[i]), GuidToString(ps_clsid), hr));
				return hr;
			}
		}

		is_registered = true;
		return S_OK;
	}

protected:
	virtual ~ComProxy()
	{
		CORE_LOG(L2, (_T("[ComProxy::~ComProxy]")));
	}

	CComPtr<IUnknown> proxy_manager_;
	CComPtr<IInternalUnknown> proxy_internal_unknown_;

	DISALLOW_COPY_AND_ASSIGN(ComProxy);
};

#pragma warning(pop)

class StdMarshalInfo : public IStdMarshalInfo
{
public:
	explicit StdMarshalInfo(bool is_machine) : is_machine_(is_machine)
	{
		CORE_LOG(L6, (_T("[StdMarshalInfo::StdMarshalInfo][%d]"), is_machine));

		VERIFY1(SUCCEEDED(ComProxy::RegisterProxyStubs()));
	}

	// IStdMarshalInfo
	STDMETHODIMP GetClassForHandler(DWORD dwDestContext, void *pvDestContext, CLSID *pClsid)
	{
		UNREFERENCED_PARAMETER(dwDestContext);
		UNREFERENCED_PARAMETER(pvDestContext);

		*pClsid = is_machine_ ? __uuidof(LongkeyComProxyMachineClass) : 
			__uuidof(LongkeyComProxyUserClass);
		return S_OK;
	}

private:
	bool is_machine_;

	DISALLOW_COPY_AND_ASSIGN(StdMarshalInfo);
};

#endif	// LONGKEY_PCAGENT_COM_PROXY_H_